import os
import requests
import configparser
import folder_paths
import time
import torchaudio
from .wavespeed_api.client import WaveSpeedClient
from .wavespeed_api.utils import tensor2images,save_audio,save_video
from .wavespeed_task_nodes import WaveSpeedTaskCreateDynamic, WaveSpeedTaskSubmit, WaveSpeedTaskStatus
from comfy.comfy_types import IO
from comfy_api.input import AudioInput, VideoInput

try:
    current_dir = os.path.dirname(os.path.abspath(__file__))
    parent_dir = os.path.dirname(current_dir)
    temp_dir = os.path.join(parent_dir, '.temp')
    if not os.path.exists(temp_dir):
        os.makedirs(temp_dir)
    config_path = os.path.join(parent_dir, 'config.ini')
    config = configparser.ConfigParser()

    if not os.path.exists(config_path):
        config['API'] = {'WAVESPEED_API_KEY': ''}
        with open(config_path, 'w') as config_file:
            config.write(config_file)

    config.read(config_path)
except Exception as e:
    print(f"Error reading or creating config file: {e}")
    config = None


class WaveSpeedAIAPIClient:
    """
    WaveSpeed AI API Client Node

    This node creates a client for connecting to the WaveSpeed AI API.
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "api_key": ("STRING", {"multiline": False, "default": ""}),
            },
        }

    RETURN_TYPES = ("WAVESPEED_AI_API_CLIENT",)
    RETURN_NAMES = ("client",)

    FUNCTION = "create_client"

    CATEGORY = "WaveSpeedAI"

    def create_client(self, api_key):
        """
        Create a WaveSpeed AI API client

        Args:
            api_key: WaveSpeed AI API key

        Returns:
            WaveSpeedAPI: WaveSpeed AI API client
        """
        wavespeed_api_key = ""
        if api_key == "":
            try:
                wavespeed_api_key = config['API']['WAVESPEED_API_KEY']
                if wavespeed_api_key == '':
                    raise ValueError('API_KEY is empty')

            except KeyError:
                raise ValueError('Unable to find API_KEY in config.ini')

        else:
            wavespeed_api_key = api_key

        return ({
            "api_key": wavespeed_api_key
        },)


class SaveAudio:
    """
    Preview Audio Node for ComfyUI
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "audio_url": ("STRING", {"forceInput": False}),
                "save_file_prefix": ("STRING", {"default": "wavespeed_audio"}),
            }
        }

    OUTPUT_NODE = True
    FUNCTION = "run"
    CATEGORY = "WaveSpeedAI"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("path",)

    def run(self, audio_url, save_file_prefix):
        if not audio_url:
            raise ValueError("No audio URL provided")

        try:
            response = requests.get(audio_url)
            response.raise_for_status()
            audio_data = response.content
        except requests.RequestException as e:
            raise RuntimeError(f"Error downloading audio: {e}")
        file_extension = os.path.splitext(audio_url)[-1]
        if not file_extension or file_extension == '.':
            file_extension = '.mp3'
        else:
            file_extension = file_extension.lower()

        full_output_folder, filename, counter, _, _ = folder_paths.get_save_image_path(save_file_prefix, folder_paths.get_output_directory())

        file = f"{filename}_{counter:05}_{file_extension}"
        file_path = os.path.join(full_output_folder, file)

        with open(file_path, "wb") as f:
            f.write(audio_data)

        return (file_path,)


class PreviewVideo:
    """
    Preview Video Node for ComfyUI

    This node allows previewing and saving videos generated by WaveSpeed AI.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "video_url": ("STRING", {"forceInput": False}),
                "save_file_prefix": ("STRING", {"default": "wavespeed_video"}),
            }
        }

    OUTPUT_NODE = True
    FUNCTION = "run"
    CATEGORY = "WaveSpeedAI"
    RETURN_TYPES = ()
    RETURN_NAMES = ()

    def run(self, video_url, save_file_prefix):
        """
        Preview and save a video

        Args:
            video_url: URL of the video
            save_file_prefix: Prefix for the saved file

        Returns:
            str: Path to the saved video file
        """
        if not video_url:
            raise ValueError("No video URL provided")

        try:
            response = requests.get(video_url)
            response.raise_for_status()
            video_data = response.content
        except requests.RequestException as e:
            raise RuntimeError(f"Error downloading video: {e}")

        file_extension = os.path.splitext(video_url)[-1]
        if not file_extension or file_extension == '.':
            file_extension = '.mp4'
        else:
            file_extension = file_extension.lower()

        output_dir = folder_paths.get_output_directory()
        filename = f"{save_file_prefix}_{int(time.time())}{file_extension}"
        file_path = os.path.join(output_dir, filename)

        with open(file_path, "wb") as f:
            f.write(video_data)

        return {"ui": {"video_url": [video_url]}, "result": ('',)}


class WanLoras:
    """
    WaveSpeed AI LoRA Parameters Node

    This node is used to set LoRA model parameters for use by other nodes.
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora1_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora1_scale": ("FLOAT", {
                    "default": 0.8,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.05,
                    "display": "number"
                }),
            },
            "optional": {
                "lora2_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora2_scale": ("FLOAT", {
                    "default": 0.8,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.05,
                    "display": "number"
                }),
                "lora3_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora3_scale": ("FLOAT", {
                    "default": 0.8,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.05,
                    "display": "number"
                }),
            }
        }

    RETURN_TYPES = ("WAVESPEED_LORAS",)
    RETURN_NAMES = ("loras",)

    FUNCTION = "create_loras"

    CATEGORY = "WaveSpeedAI"

    def create_loras(self,
                     lora1_path,
                     lora1_scale,
                     lora2_path="",
                     lora2_scale=0.8,
                     lora3_path="",
                     lora3_scale=0.8):
        """
        Create a list of LoRA parameters

        Args:
            lora1_path: Path to the first LoRA model
            lora1_scale: Weight of the first LoRA model
            lora2_path: Path to the second LoRA model
            lora2_scale: Weight of the second LoRA model
            lora3_path: Path to the third LoRA model
            lora3_scale: Weight of the third LoRA model
        Returns:
            list: List of LoRA parameters
        """
        loras = []

        # Add non-empty LoRA parameters
        for path, scale in [(lora1_path, lora1_scale),
                            (lora2_path, lora2_scale),
                            (lora3_path, lora3_scale)]:
            if path.strip():
                loras.append({"path": path.strip(), "scale": float(scale)})

        return (loras,)


class CustomLoras:
    """
    WaveSpeed AI LoRA Parameters Node

    This node is used to set LoRA model parameters for use by other nodes.
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora1_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora1_scale": ("FLOAT", {
                    "default": 0.8,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.05,
                    "display": "number"
                }),
            },
            "optional": {
                "lora2_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora2_scale": ("FLOAT", {
                    "default": 0.8,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.05,
                    "display": "number"
                }),
                "lora3_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora3_scale": ("FLOAT", {
                    "default": 0.8,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.05,
                    "display": "number"
                }),
            }
        }

    RETURN_TYPES = ("WAVESPEED_LORAS",)
    RETURN_NAMES = ("loras",)

    FUNCTION = "create_loras"

    CATEGORY = "WaveSpeedAI"

    def create_loras(self,
                     lora1_path,
                     lora1_scale,
                     lora2_path="",
                     lora2_scale=0.8,
                     lora3_path="",
                     lora3_scale=0.8):
        """
        Create a list of LoRA parameters

        Args:
            lora1_path: Path to the first LoRA model
            lora1_scale: Weight of the first LoRA model
            lora2_path: Path to the second LoRA model
            lora2_scale: Weight of the second LoRA model
            lora3_path: Path to the third LoRA model
            lora3_scale: Weight of the third LoRA model
        Returns:
            list: List of LoRA parameters
        """
        loras = []

        # Add non-empty LoRA parameters
        for path, scale in [(lora1_path, lora1_scale),
                            (lora2_path, lora2_scale),
                            (lora3_path, lora3_scale)]:
            if path.strip():
                loras.append({"path": path.strip(), "scale": float(scale)})

        return (loras,)


class FluxLoras:
    """
    WaveSpeed AI LoRA Parameter Node

    This node is used to set LoRA model parameters for use by other nodes.
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora1_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora1_scale": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 4.0,
                    "step": 0.05,
                    "display": "number"
                }),
            },
            "optional": {
                "lora2_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora2_scale": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 4.0,
                    "step": 0.05,
                    "display": "number"
                }),
                "lora3_path": ("STRING", {"multiline": False, "default": "", "tooltip": "LoRA model url ,like https://huggingface.co/username/model-name/blob/branch/model.safetensors"}),
                "lora3_scale": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 4.0,
                    "step": 0.05,
                    "display": "number"
                })
            }
        }

    RETURN_TYPES = ("WAVESPEED_LORAS",)
    RETURN_NAMES = ("loras",)

    FUNCTION = "create_loras"

    CATEGORY = "WaveSpeedAI"

    def create_loras(self,
                     lora1_path,
                     lora1_scale,
                     lora2_path="",
                     lora2_scale=1.0,
                     lora3_path="",
                     lora3_scale=1.0,):
        """
        Create a list of LoRA parameters

        Args:
            lora1_path: Path to the first LoRA model
            lora1_scale: Weight of the first LoRA model
            lora2_path: Path to the second LoRA model
            lora2_scale: Weight of the second LoRA model
            lora3_path: Path to the third LoRA model
            lora3_scale: Weight of the third LoRA model
            lora4_path: Path to the fourth LoRA model
            lora4_scale: Weight of the fourth LoRA model
            lora5_path: Path to the fifth LoRA model
            lora5_scale: Weight of the fifth LoRA model

        Returns:
            list: List of LoRA parameters
        """
        loras = []

        # Add non-empty LoRA parameters
        for path, scale in [(lora1_path, lora1_scale),
                            (lora2_path, lora2_scale),
                            (lora3_path, lora3_scale),]:
            if path.strip():
                loras.append({"path": path.strip(), "scale": float(scale)})

        return (loras,)


class UploadImage:
    """
    Upload a file to WaveSpeed AI API

    Args:
        file_path (str): Path to the file to be uploaded

    Returns:
        dict: API response containing the uploaded file information
    """
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client": ("WAVESPEED_AI_API_CLIENT",),
                "image": ("IMAGE",)
            }
        }

    DESCRIPTION = "Upload an image to WaveSpeed AI API. The link will expire soon, please do not rely on this link."

    RETURN_TYPES = ("STRING", "STRING",)
    RETURN_NAMES = ("first_image_url", "image_urls",)

    CATEGORY = "WaveSpeedAI"
    FUNCTION = "upload_file"

    def upload_file(self, client, image):
        images = tensor2images(image)
        image_urls = []
        real_client = WaveSpeedClient(api_key=client["api_key"])
        for image in images:
            image_url = real_client.upload_file(image)
            image_urls.append(image_url)
        return (image_urls[0], image_urls,)
    

class UploadVideo:
    """
    Upload a file to WaveSpeed AI API

    Args:
        file_path (str): Path to the file to be uploaded

    Returns:
        dict: API response containing the uploaded file information
    """
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client": ("WAVESPEED_AI_API_CLIENT",),
                "video": (IO.VIDEO,)
            }
        }

    DESCRIPTION = "Upload a video to WaveSpeed AI API. The link will expire soon, please do not rely on this link."

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("video_url",)

    CATEGORY = "WaveSpeedAI"
    FUNCTION = "upload_file"

    def upload_file(self, client, video:VideoInput):
        save_path = os.path.join(temp_dir, f"{int(time.time())}.mp4")
        save_video(video, save_path)
        real_client = WaveSpeedClient(api_key=client["api_key"])
        video_url = real_client.upload_file_with_type(save_path, "video/mp4")
        os.remove(save_path)
        return (video_url,)
    
class UploadAudio:
    """
    Upload a file to WaveSpeed AI API
    """
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client": ("WAVESPEED_AI_API_CLIENT",),
                "audio": (IO.AUDIO,)
            }
        }

    DESCRIPTION = "Upload an audio to WaveSpeed AI API. The link will expire soon, please do not rely on this link."

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("audio_url",)

    CATEGORY = "WaveSpeedAI"
    FUNCTION = "upload_file"

    def upload_file(self, client, audio:AudioInput):
        save_path = os.path.join(temp_dir, f"{int(time.time())}.mp3")
        save_audio(audio, save_path)
        real_client = WaveSpeedClient(api_key=client["api_key"])
        audio_url = real_client.upload_file_with_type(save_path, "audio/mp3")
        os.remove(save_path)
        return (audio_url,)

NODE_CLASS_MAPPINGS = {
    'WaveSpeedAI Client': WaveSpeedAIAPIClient,
    'WaveSpeedAI Loras Config': CustomLoras,
    'WaveSpeedAI Wan Loras': WanLoras,
    'WaveSpeedAI Flux Loras': FluxLoras,
    'WaveSpeedAI Preview Video': PreviewVideo,
    'WaveSpeedAI Save Audio': SaveAudio,
    'WaveSpeedAI Upload Image': UploadImage,
    'WaveSpeedAI Upload Video': UploadVideo,
    'WaveSpeedAI Upload Audio': UploadAudio,
    'WaveSpeedAI Task Create': WaveSpeedTaskCreateDynamic,
    'WaveSpeedAI Task Submit': WaveSpeedTaskSubmit,
    'WaveSpeedAI Task Status': WaveSpeedTaskStatus,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    'WaveSpeedAI Client': 'WaveSpeedAI Client',
    'WaveSpeedAI Loras Config': 'WaveSpeedAI Loras Config',
    'WaveSpeedAI Wan Loras': 'WaveSpeedAI Wan Loras',
    'WaveSpeedAI Flux Loras': 'WaveSpeedAI Flux and SDXL Loras',
    'WaveSpeedAI Preview Video': 'WaveSpeedAI Preview Video',
    'WaveSpeedAI Save Audio': 'WaveSpeedAI Save Audio',
    'WaveSpeedAI Upload Image': 'WaveSpeedAI Upload Image',
    'WaveSpeedAI Upload Video': 'WaveSpeedAI Upload Video',
    'WaveSpeedAI Upload Audio': 'WaveSpeedAI Upload Audio',
    'WaveSpeedAI Task Create': 'WaveSpeedAI Task Create [WIP]',
    'WaveSpeedAI Task Submit': 'WaveSpeedAI Task Submit [WIP]',
    'WaveSpeedAI Task Status': 'WaveSpeedAI Task Status [WIP]',
}
